#!/usr/bin/env python3
# core/server_core.py
# 核心服务器逻辑：任务状态管理、连接接受、任务循环、心跳监控
import socket
import threading
import time
import struct
import logging
from math import dist
import numpy as np
import datetime

from config import (
    SERVER_IP,
    SERVER_PORT,
    MAX_ROBOTS,
    ROBOT_TIMEOUT,
    HEARTBEAT_INTERVAL,
    NAV_GOAL_MAX_INTERVAL,
    SAVE_LOG_BAG,
    COMMAND_TYPES,
    VERSION,
    FRAME_SIZE,
)
from robots.robot_controller import RobotController
from policy.pursuit_policy import PursuitPolicy  # 假设已聚合
from map_manager.map_manage import MapManager
from utils.frame_builder import build_frame, parse_frame  # 工具导入
from utils.log_utils import monitor_robots_status  # 监控工具

logger = logging.getLogger(__name__)


class MultiRobotServer:
    def __init__(self, policy: PursuitPolicy, map_manager: MapManager, logger=None):
        if logger is None:
            logger = logging.getLogger(__name__)
        self.logger = logger
        self.policy = policy
        self.map_manager = map_manager
        self.logger.info("程序启动...")
        self.logger.info("开始加载地图和图文件...")
        # 机器人存储
        self.evader = None  # 逃逸者
        self.pursuers = {}  # 追捕者字典: pursuer_id -> RobotController
        self.server_socket = None
        # 任务状态
        self.game_active = False
        self.game_paused = False
        # 心跳管理
        self.heartbeat_thread = threading.Thread(
            target=self.heartbeat_loop, daemon=True
        )
        self.heartbeat_thread.start()
        # 节点变化事件
        self.node_change_event = threading.Event()
        # 连接管理锁
        self.connection_lock = threading.Lock()
        # 启动任务循环线程
        self.game_thread = threading.Thread(target=self.run_game_loop, daemon=True)
        self.game_thread.start()
        # 启动监控线程
        monitor_thread = threading.Thread(target=self._run_monitor, daemon=True)
        monitor_thread.start()
        # Logbag初始化
        if SAVE_LOG_BAG:
            self.logbag_file = open("pursuit_logbag.txt", "w")

    def start(self):
        """启动服务器监听"""
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind((SERVER_IP, SERVER_PORT))
        self.server_socket.listen(MAX_ROBOTS)
        self.logger.info(f"服务器启动，监听 {SERVER_IP}:{SERVER_PORT}")
        self.accept_connections()

    def accept_connections(self):
        """接受客户端连接"""
        while True:
            try:
                client_socket, addr = self.server_socket.accept()
                self.logger.info(f"新连接: {addr}")
                # 接收连接请求帧
                data = client_socket.recv(FRAME_SIZE)
                if len(data) != FRAME_SIZE:
                    self.logger.warning("无效连接请求长度")
                    client_socket.close()
                    continue
                # 解析帧
                frame = parse_frame(data)
                if not frame or frame["command"] != COMMAND_TYPES["ConnectionRequest"]:
                    self.logger.warning("无效的连接请求")
                    client_socket.close()
                    continue
                # 获取机器人名称
                try:
                    robot_name = frame["payload"].rstrip(b"\x00").decode("utf-8")
                except UnicodeDecodeError:
                    robot_name = "unknown_robot"
                self.logger.info(f"连接请求 | 机器人名称: {robot_name}")
                # 分配机器人ID和角色
                role, robot_id = None, None
                if self.evader is None:
                    role, robot_id = "evader", 0
                else:
                    if 1 not in self.pursuers:
                        role, robot_id = "pursuer", 1
                    elif 2 not in self.pursuers:
                        role, robot_id = "pursuer", 2
                    else:
                        self.logger.warning("拒绝连接：无可用位置")
                        client_socket.close()
                        continue
                if robot_id is None:
                    self.logger.warning("拒绝连接：无可用位置")
                    client_socket.close()
                    continue
                # 发送连接响应
                response_payload = struct.pack(">i", robot_id) + bytes(10)
                response_frame = build_frame(
                    COMMAND_TYPES["ConnectionResponse"], response_payload, robot_id
                )
                client_socket.sendall(response_frame)
                self.logger.info(f"发送连接响应，分配ID: {robot_id}, 角色: {role}")
                # 创建机器人控制器
                robot = RobotController(
                    robot_id, robot_name, role, client_socket, addr, self
                )
                if role == "evader":
                    self.evader = robot
                else:
                    self.pursuers[robot_id] = robot
                self.logger.info(f"机器人 {robot_id} ({role}) 连接成功")
                self.check_game_state()
            except Exception as e:
                self.logger.error(f"接受连接错误: {e}")
                if "client_socket" in locals() and client_socket:
                    try:
                        client_socket.close()
                    except:
                        pass

    def check_game_state(self):
        """检查任务状态并更新"""
        evader_ready = self.evader and self.evader.is_active()
        pursuer1_ready = 1 in self.pursuers and self.pursuers[1].is_active()
        pursuer2_ready = 2 in self.pursuers and self.pursuers[2].is_active()
        pursuers_ready = pursuer1_ready and pursuer2_ready
        if evader_ready and pursuers_ready:
            if not self.game_active and not self.game_paused:
                self.logger.info("所有机器人准备就绪，开始围捕!")
                self.game_active = True
                self.node_change_event.set()
                if self.evader.current_node is not None:
                    self.policy.init_belief(self.evader.current_node)
                    self.logger.info(
                        f"任务开始，信念重置到逃避者节点 {self.evader.current_node}"
                    )
        else:
            if self.game_active:
                self.logger.warning("机器人断开，暂停围捕!")
                self.game_active = False

    def run_game_loop(self):
        """任务主循环 - 事件驱动"""
        self.logger.info("[Game Loop] Game Loop 启动...")
        while True:
            self.node_change_event.wait()
            if self.game_active:
                self.calculate_and_send_strategy()
            self.node_change_event.clear()

    def send_emergency_stop_to_all(self):
        """发送紧急停止指令"""
        if self.evader and self.evader.client_socket:
            try:
                payload = bytes(14)
                frame = build_frame(
                    COMMAND_TYPES["EmergencyStop"], payload, self.evader.robot_id
                )
                self.evader.client_socket.sendall(frame)
                self.logger.info(f"[Robot-{self.evader.robot_id}] 发送紧急停止指令")
            except Exception as e:
                self.logger.error(
                    f"[Robot-{self.evader.robot_id}] 发送紧急停止失败: {e}"
                )
        for pursuer in self.pursuers.values():
            if pursuer and pursuer.client_socket:
                try:
                    payload = bytes(14)
                    frame = build_frame(
                        COMMAND_TYPES["EmergencyStop"], payload, pursuer.robot_id
                    )
                    pursuer.client_socket.sendall(frame)
                    self.logger.info(f"[Robot-{pursuer.robot_id}] 发送紧急停止指令")
                except Exception as e:
                    self.logger.error(
                        f"[Robot-{pursuer.robot_id}] 发送紧急停止失败: {e}"
                    )

    def _update_robot_nodes(self, robot):
        """手动更新机器人节点信息（基于当前位置和新地图）"""
        if robot and robot.current_position:
            x, y = robot.current_position
            new_node = self.policy.position_to_node(x, y)
            robot.current_node = new_node
            robot.target_node = new_node  # 目标设置为当前节点，避免立即移动
            self.logger.info(f"[Robot-{robot.robot_id}] 节点手动更新: {new_node}")

    def calculate_and_send_strategy(self):
        """计算并发送策略（逃避者独立计算/发送）"""
        if not self.game_active:
            return
        if not self.evader or not self.evader.current_position:
            self.logger.warning("逃逸者未连接或位置未知，跳过策略计算")
            return
        pursuer1 = self.pursuers.get(1)
        pursuer2 = self.pursuers.get(2)
        if (
            not pursuer1
            or not pursuer2
            or not pursuer1.current_position
            or not pursuer2.current_position
        ):
            self.logger.warning("追捕者未连接或位置未知，跳过策略计算")
            return
        # 获取当前节点
        evader_node = self.evader.current_node
        pursuer1_node = pursuer1.current_node
        pursuer2_node = pursuer2.current_node
        # 检查捕获
        if self.policy.is_captured(evader_node, pursuer1_node, pursuer2_node):
            self.logger.info("逃逸者已被捕获！任务结束")
            self.game_active = False
            self.policy.increment_capture_count()
            self.send_emergency_stop_to_all()
            return

        current_time = time.time()
        force_update = (
            (current_time - self.evader.last_target_send_time > NAV_GOAL_MAX_INTERVAL)
            or (current_time - pursuer1.last_target_send_time > NAV_GOAL_MAX_INTERVAL)
            or (current_time - pursuer2.last_target_send_time > NAV_GOAL_MAX_INTERVAL)
        )

        # 步骤1: 独立计算并发送逃避者移动（本地运算）
        evader_next = self.policy.calculate_evader_move(
            evader_node, pursuer1_node, pursuer2_node
        )
        self.logger.info(f"[Game Loop] 逃避者计算: {evader_node} -> {evader_next}")
        if evader_next != self.evader.target_node or force_update:
            self.evader.update_target(evader_next)
            self.evader.send_target()

        # 步骤2: 计算并发送追捕者移动（远程运算，基于当前逃避者位置）
        pursuer1_next, pursuer2_next = self.policy.calculate_pursuers_move(
            evader_node, pursuer1_node, pursuer2_node, self.policy.belief_manager
        )
        self.logger.info(
            f"[Game Loop] 追捕者计算: 追捕者1 {pursuer1_node} -> {pursuer1_next}, "
            f"追捕者2 {pursuer2_node} -> {pursuer2_next}"
        )
        if pursuer1_next != pursuer1.target_node or force_update:
            pursuer1.update_target(pursuer1_next)
            pursuer1.send_target()
        if pursuer2_next != pursuer2.target_node or force_update:
            pursuer2.update_target(pursuer2_next)
            pursuer2.send_target()

    def heartbeat_loop(self):
        """心跳循环"""
        while True:
            time.sleep(HEARTBEAT_INTERVAL)
            current_time = time.time()
            if self.evader and self.evader.status == "active":
                self.evader.send_heartbeat()
                if self.game_active and (
                    current_time - self.evader.last_target_send_time
                    > NAV_GOAL_MAX_INTERVAL
                ):
                    self.logger.info(f"[Heartbeat] 逃逸者超时，触发强制更新")
                    threading.Thread(
                        target=self.calculate_and_send_strategy, daemon=True
                    ).start()
            for pursuer in self.pursuers.values():
                if pursuer.status == "active":
                    pursuer.send_heartbeat()
                    if self.game_active and (
                        current_time - pursuer.last_target_send_time
                        > NAV_GOAL_MAX_INTERVAL
                    ):
                        self.logger.info(
                            f"[Heartbeat] 追捕者 {pursuer.robot_id} 超时，触发强制更新"
                        )
                        threading.Thread(
                            target=self.calculate_and_send_strategy, daemon=True
                        ).start()

    def _run_monitor(self):
        """运行监控（使用工具）"""
        monitor_robots_status(self)  # 委托给工具，处理日志、清理等

    def on_robot_node_changed(self):
        """节点变化事件"""
        self.node_change_event.set()

    def toggle_game_active(self):
        """切换任务状态"""
        self.game_active = not self.game_active
        if self.game_active:
            self.node_change_event.set()
        self.logger.info(f"任务状态切换: {'活跃' if self.game_active else '暂停'}")

    def position_reset(self):
        """位置重置"""
        self.game_active = False
        if self.evader:
            self.evader.current_position = None
            self.evader.current_node = None
            self.evader.target_position = None
            self.evader.target_node = None
        for pursuer in self.pursuers.values():
            pursuer.current_position = None
            pursuer.current_node = None
            pursuer.target_position = None
            pursuer.target_node = None
        self.logger.info("位置重置完成")
        # 手动更新节点信息（如果位置有效，基于新地图重新映射）
        self._update_robot_nodes(self.evader)
        for pursuer in self.pursuers.values():
            self._update_robot_nodes(pursuer)

    def policy_reload(self):
        """重新加载策略"""
        try:
            self.logger.info("开始重新加载 PursuitPolicy...")
            if self.game_active:
                self.game_active = False
                self.node_change_event.clear()
                self.logger.info("任务已暂停以便重新加载策略")
            self.policy.reset()
            self.logger.info("PursuitPolicy 重新加载成功！")
            # 手动更新机器人节点信息（基于新策略/地图）
            self._update_robot_nodes(self.evader)
            for pursuer in self.pursuers.values():
                self._update_robot_nodes(pursuer)
        except Exception as e:
            self.logger.error(f"重新加载 PursuitPolicy 失败: {e}")
            raise

    def get_all_robots_info(self):
        """获取所有机器人信息"""
        robots = []
        if self.evader:
            robots.append(self.evader.get_info())
        for pursuer in self.pursuers.values():
            robots.append(pursuer.get_info())
        return robots

    def get_game_active(self):
        return self.game_active

    def get_capture_count(self):
        return self.policy.get_capture_count()

    def get_belief_pr(self):
        return self.policy.get_belief_pr()

    def get_belief_pos(self):
        return self.policy.get_belief_pos()
